from typing import Optional
from phi.assistant import Assistant
from phi.knowledge import AssistantKnowledge
from phi.llm.ollama import Ollama
from phi.embedder.ollama import OllamaEmbedder  # 移除 OpenAI 依赖
from phi.vectordb.pgvector import PgVector2
from phi.storage.assistant.postgres import PgAssistantStorage

db_url = "postgresql+psycopg://ai:ai@localhost:5532/ai?sslmode=disable"

def get_Ollama_assistant(
    llm_model: str = "deepseek-r1:1.5b",  # 严格匹配 Ollama 模型名
    embeddings_model: str = "nomic-embed-text:latest",
    user_id: Optional[str] = None,
    run_id: Optional[str] = None,
    debug_mode: bool = True,
) -> Assistant:
    """Initialize a production-grade RAG Assistant with enhanced query handling"""

    # 确保始终使用 Ollama 嵌入模型
    embedder = OllamaEmbedder(
        model=embeddings_model,
        dimensions=768,
        base_url="http://localhost:11434"  # 显式设置 Ollama 服务地址
    )
    embeddings_table = "ollama_rag_documents"

    return Assistant(
        name="ollama_rag_assistant",
        run_id=run_id,
        user_id=user_id,
        llm=Ollama(
            model=llm_model,
            base_url="http://localhost:11434",  # 显式设置 API 端点
            temperature=0.3,
            system_prompt="严谨的知识服务系统"
        ),
        storage=PgAssistantStorage(table_name="ollama_rag_assistant", db_url=db_url),
        knowledge_base=AssistantKnowledge(
            vector_db=PgVector2(
                db_url=db_url,
                collection=embeddings_table,
                embedder=embedder,
            ),
            num_documents=3,
        ),
        description="基于Ollama语言模型和PostgreSQL向量数据库的企业级知识服务系统",
        instructions=[
            "严格依据上下文内容生成专业级响应",
            "注意检索知识库内容生成专业级响应",
            "采用分章节结构：核心结论、技术解析、数据参考",
            "使用标准文献引用格式标注来源 [1][2]",
            "保持学术论文的严谨表述，禁用口语化词汇",
            "复杂概念需附加Markdown公式说明",
            "涉及多文档时建立对比分析表格"
        ],
        add_references_to_prompt=True,
        reference_format="[来源{index}]",  # 标准化引用格式
        markdown=True,
        add_chat_history_to_messages=True,
        num_history_messages=4,
        add_datetime_to_instructions=True,
        debug_mode=debug_mode,
    )